{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cc001dde",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding PyHealth to sys.path: /home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "pyhealth_path = os.path.dirname(os.getcwd())\n",
    "if pyhealth_path not in sys.path:\n",
    "    print(f\"Adding PyHealth to sys.path: {pyhealth_path}\")\n",
    "    sys.path.insert(0, pyhealth_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9887d8db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyhealth_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3e05a68b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyhealth.tasks import BaseTask\n",
    "from typing import Any, Dict, List, Optional\n",
    "from datetime import datetime\n",
    "\n",
    "class MortalityPredictionMIMIC3Heterogeneous(BaseTask):\n",
    "    \"\"\"Task for predicting mortality using MIMIC-III dataset with text data.\n",
    "\n",
    "    This task aims to predict whether the patient will decease in the next hospital\n",
    "    visit based on clinical information from the current visit.\n",
    "    \"\"\"\n",
    "\n",
    "    task_name: str = \"MortalityPredictionMIMIC3\"\n",
    "    input_schema: Dict[str, str] = {\n",
    "        \"conditions\": \"sequence\",\n",
    "        \"procedures\": \"sequence\",\n",
    "        \"drugs\": \"sequence\",\n",
    "    }\n",
    "    output_schema: Dict[str, str] = {\"mortality\": \"binary\"}\n",
    "\n",
    "    def __call__(self, patient: Any) -> List[Dict[str, Any]]:\n",
    "        \"\"\"Processes a single patient for the mortality prediction task.\"\"\"\n",
    "        samples = []\n",
    "\n",
    "        # We will drop the last visit\n",
    "        visits = patient.get_events(event_type=\"admissions\")\n",
    "\n",
    "        if len(visits) <= 1:\n",
    "            return []\n",
    "\n",
    "        for i in range(len(visits) - 1):\n",
    "            visit = visits[i]\n",
    "            next_visit = visits[i + 1]\n",
    "\n",
    "            # Check discharge status for mortality label - more robust handling\n",
    "            if next_visit.hospital_expire_flag not in [0, 1, \"0\", \"1\"]:\n",
    "                mortality_label = 0\n",
    "            else:\n",
    "                mortality_label = int(next_visit.hospital_expire_flag)\n",
    "\n",
    "            # Convert string timestamps to datetime objects\n",
    "            try:\n",
    "                # Check the type and convert if necessary\n",
    "                if isinstance(visit.dischtime, str):\n",
    "                    discharge_time = datetime.strptime(\n",
    "                        visit.dischtime, \"%Y-%m-%d %H:%M:%S\"\n",
    "                    )\n",
    "                else:\n",
    "                    discharge_time = visit.dischtime\n",
    "            except (ValueError, AttributeError):\n",
    "                # If conversion fails, skip this visit\n",
    "                print(\"Error parsing discharge time:\", visit.dischtime)\n",
    "                continue\n",
    "\n",
    "            # Get clinical codes\n",
    "            diagnoses = patient.get_events(\n",
    "                event_type=\"diagnoses_icd\",\n",
    "                start=visit.timestamp,\n",
    "                end=discharge_time,  # Now using a datetime object\n",
    "            )\n",
    "            procedures = patient.get_events(\n",
    "                event_type=\"procedures_icd\",\n",
    "                start=visit.timestamp,\n",
    "                end=discharge_time,  # Now using a datetime object\n",
    "            )\n",
    "            prescriptions = patient.get_events(\n",
    "                event_type=\"prescriptions\",\n",
    "                start=visit.timestamp,\n",
    "                end=discharge_time,  # Now using a datetime object\n",
    "            )\n",
    "\n",
    "            conditions = [event.icd9_code for event in diagnoses]\n",
    "            procedures_list = [event.icd9_code for event in procedures]\n",
    "            drugs = [event.drug for event in prescriptions]\n",
    "\n",
    "            # Exclude visits without condition, procedure, or drug code\n",
    "            samples.append(\n",
    "                {\n",
    "                    \"hadm_id\": visit.hadm_id,\n",
    "                    \"patient_id\": patient.patient_id,\n",
    "                    \"conditions\": conditions,\n",
    "                    \"procedures\": procedures_list,\n",
    "                    \"drugs\": drugs,\n",
    "                    \"mortality\": mortality_label,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        return samples\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da8360f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No config path provided, using default config\n",
      "Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: False)\n",
      "Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth/pyhealth/datasets/mimic3.py:50: UserWarning: Events from prescriptions table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some column names were converted to lowercase\n",
      "Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
      "Some column names were converted to lowercase\n",
      "Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n",
      "Some column names were converted to lowercase\n",
      "Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n",
      "Some column names were converted to lowercase\n",
      "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
      "Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n",
      "Some column names were converted to lowercase\n",
      "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
      "Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n",
      "Some column names were converted to lowercase\n",
      "Scanning table: noteevents from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv.gz\n",
      "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv\n",
      "Some column names were converted to lowercase\n",
      "Preprocessing table: noteevents with preprocess_noteevents\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import MIMIC3Dataset\n",
    "dataset = MIMIC3Dataset(\n",
    "    root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n",
    "    tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\", \"noteevents\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dd5c282f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting task MortalityPredictionMIMIC3 for mimic3 base dataset...\n",
      "Loading cached samples from cache/MortalityPredictionMIMIC3.parquet\n",
      "Loaded 2776 cached samples\n",
      "Label mortality vocab: {0: 0, 1: 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing samples: 100%|██████████| 2776/2776 [00:00<00:00, 63155.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 2776 samples for task MortalityPredictionMIMIC3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.tasks.mortality_prediction import MortalityPredictionMIMIC3\n",
    "from pyhealth.datasets import split_by_patient, get_dataloader\n",
    "mimic3_mortality_prediction = MortalityPredictionMIMIC3Heterogeneous()\n",
    "samples = dataset.set_task(mimic3_mortality_prediction, num_workers=1, cache_dir=\"cache/\") # use default task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "06b5a9a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyhealth.datasets import split_by_sample\n",
    "\n",
    "\n",
    "train_dataset, val_dataset, test_dataset = split_by_sample(\n",
    "    dataset=samples,\n",
    "    ratios=[0.7, 0.1, 0.2]\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "medical_coding_demo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
